package com.dynamic.mybatis.core.session;

import com.baomidou.mybatisplus.core.toolkit.ArrayUtils;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.dynamic.mybatis.core.datasource.DataSourceProperty;
import com.dynamic.mybatis.core.datasource.DynamicDataSourceHolder;
import com.dynamic.mybatis.core.datasource.DynamicMultipleDataSource;
import com.dynamic.mybatis.core.intercept.DynamicDataSourceHandler;
import com.dynamic.mybatis.core.metadata.DynamicConfiguration;
import com.dynamic.mybatis.core.metadata.DynamicMappedStatement;
import com.dynamic.mybatis.core.utils.StringUtils;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.SQLException;
import java.util.*;

/**
 * @Description
 * @Author xs
 * @Date 2023/3/15 9:23
 */
public abstract class DynamicSqlSessionDataSource {

    Set<DynamicDataSourceHandler> dynamicDataSourceHandlers = new HashSet<>();

    public abstract SqlSession getSqlSession();

    protected abstract SqlSessionFactory getSqlSqlSessionFactory();

    protected abstract DynamicSqlSession getDynamicSqlSession();


    public Set<DynamicDataSourceHandler> getDynamicDataSourceHandlers() {
        return dynamicDataSourceHandlers;
    }


    public DataSource getDataSource() {
        DataSource dataSource = getSqlSession().getConfiguration().getEnvironment().getDataSource();
        return dataSource;
    }

    public DataSource getDataSource(String dataSourceCode) {
        Configuration configuration = getSqlSession().getConfiguration();
        if (configuration == null || !(configuration instanceof DynamicConfiguration)) {
            return getDataSource();
        }
        DataSource dataSource = getDataSource();
        if (!(dataSource instanceof DynamicMultipleDataSource)) {
            return dataSource;
        }
        DynamicMultipleDataSource dynamicMultipleDataSource = (DynamicMultipleDataSource) dataSource;
        return dynamicMultipleDataSource.getDataSource(dataSourceCode);
    }

    public DataSourceProperty getDataSourceProperty(String dataSourceCode) {
        Configuration configuration = getSqlSession().getConfiguration();
        if (configuration == null || !(configuration instanceof DynamicConfiguration)) {
            throw new RuntimeException("configuration is null or configuration is not  instanceof DynamicConfiguration");
        }
        DataSource dataSource = getDataSource();
        if (!(dataSource instanceof DynamicMultipleDataSource)) {
            throw new RuntimeException("main dataSource  is not instanceof DynamicMultipleDataSource");
        }
        DynamicMultipleDataSource dynamicMultipleDataSource = (DynamicMultipleDataSource) dataSource;
        return dynamicMultipleDataSource.getDataSourceProperty(dataSourceCode);
    }


    public void addDataSourceHandler(Collection<DynamicDataSourceHandler> dataSourceHandlers) {
        for (DynamicDataSourceHandler intercept : dataSourceHandlers) {
            addDataSourceHandler(intercept);
        }
    }

    public void addDataSourceHandler(DynamicDataSourceHandler dataSourceHandler) {
        if (dataSourceHandler == null) {
            return;
        }
        this.dynamicDataSourceHandlers.add(dataSourceHandler);
    }


    //如果数据源是连接池源，则返回的connection不要关闭
    public Connection getConnection(String mappedStatementId) {
        Configuration configuration = getSqlSession().getConfiguration();
        if (configuration == null || !(configuration instanceof DynamicConfiguration)) {
            return getConnection();
        }
        DataSource dataSource = getDataSource();
        if (!(dataSource instanceof DynamicMultipleDataSource)) {
            return getConnection(dataSource, null);
        }
        DynamicConfiguration dc = (DynamicConfiguration) configuration;
        DynamicMappedStatement dms = dc.getDynamicMappedStatement(mappedStatementId);
        if (dms == null) {
            return getConnection(dataSource, null);
        }
        Object datasourceKey;
        datasourceKey = dms.getTargetDataSourceKey();
        if (datasourceKey == null) {
            if (ArrayUtils.isNotEmpty(dms.getDynamicDataSourceKeys())) {
                datasourceKey = dms.getDynamicDataSourceKeys()[0];
            } else {
                return getConnection(dataSource, null);
            }
        }
        if (datasourceKey == null) {
            return getConnection(dataSource, null);
        }
        return getConnection(dataSource, datasourceKey);
    }

    public Connection getConnectionByDataSourceKey(String dataSourceKey) {
        Configuration configuration = getSqlSession().getConfiguration();
        if (configuration == null || !(configuration instanceof DynamicConfiguration)) {
            return getSqlSession().getConnection();
        }
        DataSource dataSource = configuration.getEnvironment().getDataSource();
        if (!(dataSource instanceof AbstractRoutingDataSource)) {
            return getSqlSession().getConnection();
        }
        return getConnection(dataSource, dataSourceKey);
    }

    public void addDataSource(String dataSourceKey, DataSource dataSource) {
        addDataSource(dataSourceKey, dataSource, null, null);
    }

    /**
     * 添加数据源
     *
     * @param dataSourceKey 数据源key
     * @param dataSource    数据源
     * @param attributes    通用的属性key com.dynamic.mybatis.core.datasource.DataSourceAttributesKey或者自定义属性key
     */

    public void addDataSource(String dataSourceKey, DataSource dataSource, DataSourceProperty property, Map<Object, Object> attributes) {
        addDataSource(dataSourceKey, dataSource, property, attributes, false);
    }

    /**
     * 添加数据源
     *
     * @param dataSourceKey 数据源key
     * @param dataSource    数据源
     * @param attributes    通用的属性key com.dynamic.mybatis.core.datasource.DataSourceAttributesKey或者自定义属性key
     * @param handler       是否调用 com.dynamic.mybatis.core.intercept.DynamicDataSourceHandler 实现类
     */

    public void addDataSource(String dataSourceKey, DataSource dataSource, DataSourceProperty property, Map<Object, Object> attributes, boolean handler) {
        DataSource originalDataSource = getSqlSession().getConfiguration().getEnvironment().getDataSource();
        if (!(originalDataSource instanceof DynamicMultipleDataSource)) {
            throw new RuntimeException("mybatis configuration environment initialized dataSource is not multipleDataSource.");
        }
        DynamicMultipleDataSource dynamicMultipleDataSource = (DynamicMultipleDataSource) originalDataSource;


        DynamicDataSourceHandler.OPERATE operate = dynamicMultipleDataSource.hasDataSource(dataSourceKey) ?
                DynamicDataSourceHandler.OPERATE.UPDATE : DynamicDataSourceHandler.OPERATE.ADD;
        //处理
        if (handler) {
            beforeHandle(operate, dataSourceKey, dataSource, property, attributes);
        }

        dynamicMultipleDataSource.put(dataSourceKey, dataSource);
        dynamicMultipleDataSource.setProperty(property);
        dynamicMultipleDataSource.setAttributes(dataSourceKey, attributes);
        afterDataSourcePropertiesSet();

        if (handler) {
            afterHandle(operate, dataSourceKey, dataSource, property, attributes);
        }
    }


    public void afterDataSourcePropertiesSet() {
        DataSource originalDataSource = getSqlSession().getConfiguration().getEnvironment().getDataSource();
        if (!(originalDataSource instanceof DynamicMultipleDataSource)) {
            return;
        }
        DynamicMultipleDataSource dynamicMultipleDataSource = (DynamicMultipleDataSource) originalDataSource;
        Map<Object, DataSource> dataSourceMap = dynamicMultipleDataSource.getDataSources();
        if (CollectionUtils.isEmpty(dataSourceMap)) {
            return;
        }
        for (Map.Entry<Object, DataSource> entry : dataSourceMap.entrySet()) {

            DataSourceProperty property = dynamicMultipleDataSource.getProperty(entry.getKey());
            property = (property == null ? getDataSourceProperty(entry.getValue()) : property);

            if (StringUtils.isBlank(property.getCode())) {
                String key = String.valueOf(entry.getKey());
                int index = key.lastIndexOf(".");
                property.setTenantCode(key.substring(0, index));
                property.setCode(key.substring(index + 1));
            }

            if (StringUtils.isEmpty(property.getDatabaseType()) && StringUtils.isNotBlank(property.getUrl())) {
                String[] parts = property.getUrl().split("/|:|\\?");
                if (parts.length > 1) {
                    property.setDatabaseType(parts[1]);
                }
                if (parts.length > 6) {
                    property.setDatabaseName(parts[6]);
                }
            }
            dynamicMultipleDataSource.setProperty(property);
        }
    }

    public DataSourceProperty getDataSourceProperty(DataSource dataSource) {
        DataSourceProperty property = new DataSourceProperty();
        try (Connection connection = dataSource.getConnection();) {
            DatabaseMetaData databaseMetaData = connection.getMetaData();
            property.setUrl(databaseMetaData.getURL());
            String username = databaseMetaData.getUserName();
            if (StringUtils.isNotBlank(username)) {
                property.setUsername(username.split("@")[0]);
            }
            property.setDriverVersion(databaseMetaData.getDriverVersion());
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return property;
    }

    /**
     * 　* @Description: 设置主数据源
     *
     */
    public void setMainDataSource(String dataSourceKey) {
        DataSource originalDataSource = getSqlSession().getConfiguration().getEnvironment().getDataSource();
        if ((originalDataSource instanceof DynamicMultipleDataSource)) {
            throw new RuntimeException("mybatis configuration environment initialized dataSource is not multipleDataSource.");
        }
        DynamicMultipleDataSource multipleDataSource = (DynamicMultipleDataSource) originalDataSource;
        multipleDataSource.setMainDataSource(dataSourceKey);
    }


    //执行前拦截
    public void beforeHandle(DynamicDataSourceHandler.OPERATE operate, Object dataSourceKey, DataSource dataSource, DataSourceProperty property, Map<Object, Object> attributes) {
        for (DynamicDataSourceHandler handler : this.dynamicDataSourceHandlers) {
            handler.after(operate, dataSourceKey, dataSource, property, attributes);
        }
    }

    //执行前拦截
    public void afterHandle(DynamicDataSourceHandler.OPERATE operate, Object dataSourceKey, DataSource dataSource, DataSourceProperty property, Map<Object, Object> attributes) {
        for (DynamicDataSourceHandler handler : this.dynamicDataSourceHandlers) {
            handler.after(operate, dataSourceKey, dataSource, property, attributes);
        }
    }

    public Map<Object, Object> getDataSourceAttributes(String dataSourceKey) {
        Configuration configuration = getSqlSession().getConfiguration();
        if (configuration == null || !(configuration instanceof DynamicConfiguration)) {
            return Collections.emptyMap();
        }
        DataSource dataSource = configuration.getEnvironment().getDataSource();
        if (!(dataSource instanceof DynamicMultipleDataSource)) {
            return Collections.emptyMap();
        }
        return ((DynamicMultipleDataSource) dataSource).getAttributes(dataSourceKey);
    }

    public Object getTargetDataSourceKey(DynamicMappedStatement dms, Object dataSourceKey) {
        Object targetDataSourceKey = null;
        if (!Objects.isNull(dataSourceKey)) {
            targetDataSourceKey = dataSourceKey;
        } else if (!Objects.isNull(dms.getTargetDataSourceKey())) {
            targetDataSourceKey = dms.getTargetDataSourceKey();
        } else if (ArrayUtils.isNotEmpty(dms.getDynamicDataSourceKeys())) {
            targetDataSourceKey = dms.getDynamicDataSourceKeys()[0];
        }
        return targetDataSourceKey;
    }

    //从非连接池的数据源获取链接，如果非费用使用完以后要关闭
    private Connection getConnection() {
        try {
            return getDataSource().getConnection();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private Connection getConnection(DataSource dataSource, Object dataSourceKey) {
        if (!(dataSource instanceof DynamicMultipleDataSource)) {
            return null;
        }
        try {
            DynamicDataSourceHolder.set(dataSourceKey);
            return dataSource.getConnection();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            DynamicDataSourceHolder.clear();
        }
    }
}